import numpy as np
import struct
import os

def im2col(input_data, kernel_size, stride=1, padding=0):
    input_height, input_width, channels = input_data.shape
    out_height = (input_height + 2 * padding - kernel_size) // stride + 1
    out_width  = (input_width  + 2 * padding - kernel_size) // stride + 1
    output_data = np.zeros((out_height * out_width, kernel_size * kernel_size * channels), dtype=np.int32)

    for oh in range(out_height):
        for ow in range(out_width):
            for c in range(channels):
                for kh in range(kernel_size):
                    for kw in range(kernel_size):
                        ih = oh * stride - padding + kh
                        iw = ow * stride - padding + kw
                        if 0 <= ih < input_height and 0 <= iw < input_width:
                            output_data[oh * out_width + ow, c * kernel_size * kernel_size + kh * kernel_size + kw] = input_data[ih, iw, c]
                        else:
                            output_data[oh * out_width + ow, c * kernel_size * kernel_size + kh * kernel_size + kw] = 0
    return output_data.reshape(out_height * out_width, -1)

def convolution(input_data, weight, bias=None, stride=1, padding=0):
    input_height, input_width, in_channels = input_data.shape
    out_channels, _, kernel_height, kernel_width = weight.shape

    out_height = (input_height + 2 * padding - kernel_height) // stride + 1
    out_width = (input_width + 2 * padding - kernel_width) // stride + 1
    output = np.zeros((out_height, out_width, out_channels), dtype=np.int32)

    for oc in range(out_channels):
        for oh in range(out_height):
            for ow in range(out_width):
                conv_sum = 0
                for ic in range(in_channels):
                    for kh in range(kernel_height):
                        for kw in range(kernel_width):
                            ih = oh * stride - padding + kh
                            iw = ow * stride - padding + kw
                            if 0 <= ih < input_height and 0 <= iw < input_width:
                                conv_sum += int(input_data[ih, iw, ic]) * int(weight[oc, ic, kh, kw])
                if bias is not None:
                    conv_sum += bias[oc]
                output[oh, ow, oc] = conv_sum
    return output

def convolution_with_im2col(input_data, weight, bias=None, stride=1, padding=0):
    input_height, input_width, in_channels = input_data.shape
    out_channels, _, kernel_height, kernel_width = weight.shape

    im2col_matrix = im2col(input_data, kernel_height, stride, padding)
    print(f"im2col_matrix: {im2col_matrix}")

    # Reshape weight for matrix multiplication
    weight_matrix = weight.reshape(out_channels, -1)
    print(f"weight_matrix: {weight_matrix}")

    # Perform convolution as matrix multiplication
    out_height = (input_height + 2 * padding - kernel_height) // stride + 1
    out_width = (input_width + 2 * padding - kernel_width) // stride + 1
    
    output = im2col_matrix @ weight_matrix.T
    
    # Add bias if provided
    if bias is not None:
        output += bias
    
    # Reshape to final output format
    return output.reshape(out_height, out_width, out_channels)


def reshape_weight_for_hardware(weight, P_OCH, P_ICH):
    """
    Reshape weight to match hardware format:
    ap_uint<P_OCH * P_ICH * W_BIT> weight[N_OCH / P_OCH][N_ICH / P_ICH][K * K]
    
    Input weight shape: (N_OCH, N_ICH, K, K)
    Output shape: (N_OCH/P_OCH, N_ICH/P_ICH, K*K, P_OCH, P_ICH)
    """
    N_OCH, N_ICH, K, _ = weight.shape
    FOLD_O = N_OCH // P_OCH
    FOLD_I = N_ICH // P_ICH
    KK = K * K
    
    # Reshape: (N_OCH, N_ICH, K, K) -> (FOLD_O, P_OCH, FOLD_I, P_ICH, K, K)
    weight_reshaped = weight.reshape(FOLD_O, P_OCH, FOLD_I, P_ICH, K, K)
    
    # Transpose to: (FOLD_O, FOLD_I, K, K, P_OCH, P_ICH)
    weight_reshaped = weight_reshaped.transpose(0, 2, 4, 5, 1, 3)
    
    # Reshape to: (FOLD_O, FOLD_I, K*K, P_OCH, P_ICH)
    weight_reshaped = weight_reshaped.reshape(FOLD_O, FOLD_I, KK, P_OCH, P_ICH)
    
    return weight_reshaped


def save_input_to_binary(input_data, filename, P_ICH, A_BIT=8):
    """
    Save input data to binary file for SystemVerilog testbench
    Input shape: (height, width, channels)
    Saved as packed vectors: each vector contains P_ICH channels
    """
    height, width, channels = input_data.shape
    
    # Convert to flat format: (H*W, channels)
    input_flat = input_data.reshape(height * width, channels)
    
    with open(filename, 'wb') as f:
        for hw in range(height * width):
            for ch_start in range(0, channels, P_ICH):
                # Pack P_ICH channels into one vector
                vec = input_flat[hw, ch_start:ch_start+P_ICH]
                for val in vec:
                    # Write as unsigned 8-bit integer
                    f.write(struct.pack('B', int(val) & 0xFF))
    
    print(f"Saved input (binary) to {filename}")
    print(f"  Shape: H*W={height*width}, channels={channels}")
    print(f"  Vectors: {height * width * (channels // P_ICH)}")


def save_input_to_text(input_data, filename):
    """
    Save input data to text file
    """
    height, width, channels = input_data.shape
    
    with open(filename, 'w') as f:
        f.write(f"# Input shape: height={height}, width={width}, channels={channels}\n")
        for h in range(height):
            for w in range(width):
                for c in range(channels):
                    val = input_data[h, w, c]
                    f.write(f"{int(val)}\n")
    
    print(f"Saved input (text) to {filename}")


def save_input_to_dat(input_data, filename, P_ICH, A_BIT=8):
    """
    Save input data to .dat file for $readmemh
    Input shape: (height, width, channels)
    Each line contains P_ICH packed values in hex
    """
    height, width, channels = input_data.shape
    
    with open(filename, 'w') as f:
        for h in range(height):
            for w in range(width):
                for ch_start in range(0, channels, P_ICH):
                    # Pack P_ICH channels into one hex line
                    hex_values = []
                    for ch in range(ch_start, min(ch_start + P_ICH, channels)):
                        val = int(input_data[h, w, ch])
                        hex_values.append(f"{val:02x}")
                    # Write in little-endian order (reverse)
                    f.write(''.join(reversed(hex_values)) + '\n')
    
    print(f"Saved input (dat) to {filename}")
    print(f"  Format: hex values for $readmemh")


def save_im2col_to_binary(im2col_data, filename, P_ICH, A_BIT=8):
    """
    Save im2col data to binary file
    im2col_data shape: (N_vectors, channels*K*K)
    """
    n_vec, total_ch = im2col_data.shape
    
    with open(filename, 'wb') as f:
        for vec_idx in range(n_vec):
            for ch_start in range(0, total_ch, P_ICH):
                # Pack P_ICH values into one vector
                vec = im2col_data[vec_idx, ch_start:ch_start+P_ICH]
                for val in vec:
                    # Write as unsigned 8-bit integer
                    f.write(struct.pack('B', int(val) & 0xFF))
    
    print(f"Saved im2col (binary) to {filename}")
    print(f"  Shape: N_vectors={n_vec}, total_channels={total_ch}")


def save_im2col_to_text(im2col_data, filename):
    """
    Save im2col data to text file
    """
    n_vec, total_ch = im2col_data.shape
    
    with open(filename, 'w') as f:
        f.write(f"# Im2col shape: N_vectors={n_vec}, total_channels={total_ch}\n")
        for vec_idx in range(n_vec):
            for ch in range(total_ch):
                val = im2col_data[vec_idx, ch]
                f.write(f"{int(val)}\n")
    
    print(f"Saved im2col (text) to {filename}")


def save_im2col_to_dat(im2col_data, filename, P_ICH, A_BIT=8):
    """
    Save im2col data to .dat file for $readmemh
    im2col_data shape: (N_vectors, channels*K*K)
    Each line contains P_ICH packed values in hex
    """
    n_vec, total_ch = im2col_data.shape
    
    with open(filename, 'w') as f:
        for vec_idx in range(n_vec):
            for ch_start in range(0, total_ch, P_ICH):
                # Pack P_ICH values into one hex line
                hex_values = []
                for ch in range(ch_start, min(ch_start + P_ICH, total_ch)):
                    val = int(im2col_data[vec_idx, ch])
                    hex_values.append(f"{val:02x}")
                # Write in little-endian order (reverse)
                f.write(''.join(reversed(hex_values)) + '\n')
    
    print(f"Saved im2col (dat) to {filename}")
    print(f"  Format: hex values for $readmemh")


def save_weight_to_binary(weight, filename, P_OCH, P_ICH, W_BIT=8):
    """
    Save weight data to binary file for SystemVerilog testbench
    Weight shape after reshape: (FOLD_O, FOLD_I, K*K, P_OCH, P_ICH)
    """
    FOLD_O, FOLD_I, KK, _, _ = weight.shape
    
    with open(filename, 'wb') as f:
        for fo in range(FOLD_O):
            for fi in range(FOLD_I):
                for kk in range(KK):
                    # Pack in order: [och0_ich0, och0_ich1, ..., och0_ichN, och1_ich0, ...]
                    for o in range(P_OCH):
                        for i in range(P_ICH):
                            val = weight[fo, fi, kk, o, i]
                            # Write as signed 8-bit integer
                            f.write(struct.pack('b', int(val)))
    
    print(f"Saved weight (binary) to {filename}")
    print(f"  Shape: FOLD_O={FOLD_O}, FOLD_I={FOLD_I}, K*K={KK}")
    print(f"  Total entries: {FOLD_O * FOLD_I * KK}")


def save_weight_to_text(weight, filename):
    """
    Save weight data to text file
    Weight shape: (FOLD_O, FOLD_I, K*K, P_OCH, P_ICH)
    """
    FOLD_O, FOLD_I, KK, P_OCH, P_ICH = weight.shape
    
    with open(filename, 'w') as f:
        f.write(f"# Weight shape: FOLD_O={FOLD_O}, FOLD_I={FOLD_I}, K*K={KK}, P_OCH={P_OCH}, P_ICH={P_ICH}\n")
        for fo in range(FOLD_O):
            for fi in range(FOLD_I):
                for kk in range(KK):
                    for o in range(P_OCH):
                        for i in range(P_ICH):
                            val = weight[fo, fi, kk, o, i]
                            f.write(f"{int(val)}\n")
    
    print(f"Saved weight (text) to {filename}")


def save_weight_to_dat(weight, filename, P_OCH, P_ICH, W_BIT=8):
    """
    Save weight data to .dat file for $readmemh
    Weight shape: (FOLD_O, FOLD_I, K*K, P_OCH, P_ICH)
    Each line contains one complete weight vector: P_OCH*P_ICH weights packed together
    """
    FOLD_O, FOLD_I, KK, _, _ = weight.shape
    
    with open(filename, 'w') as f:
        for fo in range(FOLD_O):
            for fi in range(FOLD_I):
                for kk in range(KK):
                    # Pack P_OCH*P_ICH weights into one hex line
                    hex_values = []
                    for o in range(P_OCH):
                        for i in range(P_ICH):
                            val = int(weight[fo, fi, kk, o, i])
                            # Convert to unsigned representation for hex
                            if val < 0:
                                val = val & ((1 << W_BIT) - 1)
                            hex_values.append(f"{val:02x}")
                    # Write in little-endian order (reverse for readmemh)
                    f.write(''.join(reversed(hex_values)) + '\n')
    
    print(f"Saved weight (dat) to {filename}")
    print(f"  Format: hex values for $readmemh")


def save_output_to_binary(output, filename, B_BIT=32):
    """
    Save output data to binary file
    """
    height, width, channels = output.shape
    
    with open(filename, 'wb') as f:
        for h in range(height):
            for w in range(width):
                for c in range(channels):
                    val = output[h, w, c]
                    # Write as signed 32-bit integer
                    f.write(struct.pack('i', int(val)))
    
    print(f"Saved output (binary) to {filename}")
    print(f"  Shape: height={height}, width={width}, channels={channels}")


def save_output_to_text(output, filename):
    """
    Save output data to text file
    """
    height, width, channels = output.shape
    
    with open(filename, 'w') as f:
        f.write(f"# Output shape: height={height}, width={width}, channels={channels}\n")
        for h in range(height):
            for w in range(width):
                for c in range(channels):
                    val = output[h, w, c]
                    f.write(f"{int(val)}\n")
    
    print(f"Saved output (text) to {filename}")


def save_output_to_dat(output, filename, B_BIT=32):
    """
    Save output data to .dat file for $readmemh
    Output shape: (height, width, channels)
    Each line contains one output value in hex (B_BIT bits)
    """
    height, width, channels = output.shape
    
    with open(filename, 'w') as f:
        for h in range(height):
            for w in range(width):
                for c in range(channels):
                    val = int(output[h, w, c])
                    # Convert to unsigned representation for hex
                    if val < 0:
                        val = val & ((1 << B_BIT) - 1)
                    # Format as hex with appropriate width (B_BIT/4 hex digits)
                    hex_width = B_BIT // 4
                    f.write(f"{val:0{hex_width}x}\n")
    
    print(f"Saved output (dat) to {filename}")
    print(f"  Format: hex values for $readmemh")


if __name__ == "__main__":
    # Configuration
    height, width, in_channels = 5, 5, 1
    out_channels = 1
    P_ICH, P_OCH = 1, 1
    kernel_size, stride, padding = 3, 1, 1
    A_BIT, W_BIT, B_BIT = 8, 8, 32

    # Generate test data using arange (0-127 range)
    # Input data
    input_size = height * width * in_channels
    input_data = (np.arange(input_size) % 128).astype(np.uint8)  # 0-127, uint8
    input_data = input_data.reshape(height, width, in_channels)
    
    # Weight data
    weight_size = out_channels * in_channels * kernel_size * kernel_size
    weight_data = (np.arange(weight_size) % 128).astype(np.uint8)  # 0-127, uint8
    weight = weight_data.reshape(out_channels, in_channels, kernel_size, kernel_size)
    
    bias = np.zeros(out_channels, dtype=np.int32)

    print("="*60)
    print("Configuration:")
    print(f"  Input: {height}x{width}x{in_channels}")
    print(f"  Weight: {out_channels}x{in_channels}x{kernel_size}x{kernel_size}")
    print(f"  P_ICH={P_ICH}, P_OCH={P_OCH}")
    print(f"  Kernel={kernel_size}, Stride={stride}, Padding={padding}")
    print("="*60)

    # Run convolution
    output_im2col = convolution_with_im2col(input_data, weight, bias, stride, padding)
    output_golden = convolution(input_data, weight, bias, stride, padding)

    max_diff = np.max(np.abs(output_im2col - output_golden))
    mean_diff = np.mean(np.abs(output_im2col - output_golden))

    print(f"\nOutput shape: {output_im2col.shape}")
    print(f"Max difference: {max_diff:.6f}")
    print(f"Mean difference: {mean_diff:.6f}")
    print(f"Results match: {np.allclose(output_im2col, output_golden, rtol=1e-5, atol=1e-6)}")

    # Reshape weight for hardware
    print("\n" + "="*60)
    print("Reshaping weight for hardware...")
    weight_hw = reshape_weight_for_hardware(weight, P_OCH, P_ICH)
    print(f"Weight hardware shape: {weight_hw.shape}")
    print(f"  (FOLD_O, FOLD_I, K*K, P_OCH, P_ICH) = {weight_hw.shape}")

    # Create data directory if it doesn't exist
    os.makedirs('.', exist_ok=True)

    # Save to binary files
    print("\n" + "="*60)
    print("Saving data to binary and text files...")
    
    # Save parameters to SystemVerilog header file
    with open('conv_config.svh', 'w') as f:
        f.write("// Auto-generated configuration file for conv testbench\n")
        f.write("// Generated by gen_conv.py\n\n")
        f.write(f"localparam int unsigned P_ICH = {P_ICH};\n")
        f.write(f"localparam int unsigned P_OCH = {P_OCH};\n")
        f.write(f"localparam int unsigned N_ICH = {in_channels};\n")
        f.write(f"localparam int unsigned N_OCH = {out_channels};\n")
        f.write(f"localparam int unsigned K = {kernel_size};\n")
        f.write(f"localparam int unsigned A_BIT = {A_BIT};\n")
        f.write(f"localparam int unsigned W_BIT = {W_BIT};\n")
        f.write(f"localparam int unsigned B_BIT = {B_BIT};\n")
        f.write(f"localparam int unsigned N_HW = {output_im2col.shape[1] * output_im2col.shape[2]};\n")
        f.write(f"\n// Derived parameters\n")
        f.write(f"localparam int unsigned FOLD_I = N_ICH / P_ICH;\n")
        f.write(f"localparam int unsigned FOLD_O = N_OCH / P_OCH;\n")
        f.write(f"localparam int unsigned KK = K * K;\n")
        f.write(f"localparam int unsigned WEIGHT_DEPTH = FOLD_O * FOLD_I * KK;\n")
    print("Saved configuration to conv_config.svh")
    
    # Save input
    save_input_to_binary(input_data, 'conv_input.bin', P_ICH, A_BIT)
    save_input_to_text(input_data, 'conv_input.txt')
    save_input_to_dat(input_data, 'conv_input.dat', P_ICH, A_BIT)
    
    # Generate and save im2col data
    im2col_data = im2col(input_data, kernel_size, stride, padding)
    save_im2col_to_binary(im2col_data, 'conv_im2col.bin', P_ICH, A_BIT)
    save_im2col_to_text(im2col_data, 'conv_im2col.txt')
    save_im2col_to_dat(im2col_data, 'conv_im2col.dat', P_ICH, A_BIT)
    
    # Save weight
    save_weight_to_binary(weight_hw, 'conv_weight.bin', P_OCH, P_ICH, W_BIT)
    save_weight_to_text(weight_hw, 'conv_weight.txt')
    save_weight_to_dat(weight_hw, 'conv_weight.dat', P_OCH, P_ICH, W_BIT)
    
    # Save output
    save_output_to_binary(output_golden, 'conv_output.bin', B_BIT)
    save_output_to_text(output_golden, 'conv_output.txt')
    save_output_to_dat(output_golden, 'conv_output.dat', B_BIT)